{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vwELCooy4ljr"
      },
      "source": [
        "# Prompt templates and task chains\n",
        "\n",
        "txtai has long had support for workflows. Workflows connect the input and outputs of machine learning models together to create powerful transformation and processing functions.\n",
        "\n",
        "There has been a recent surge in interest in \"model prompting\", which is the process of building a natural language description of a task and passing it to a large language model (LLM). txtai has recently improved support for task templating, which builds string outputs from a set of parameters.\n",
        "\n",
        "This notebook demonstrates how txtai workflows can be used to apply prompt templates and chain those tasks together."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ew7orE2O441o"
      },
      "source": [
        "# Install dependencies\n",
        "\n",
        "Install `txtai` and all dependencies."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LPQTb25tASIG"
      },
      "source": [
        "%%capture\n",
        "!pip install git+https://github.com/neuml/txtai#egg=txtai[api]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_YnqorRKAbLu"
      },
      "source": [
        "# Prompt workflow\n",
        "\n",
        "First, we'll look at building a workflow with a series of model prompts. This workflow creates a conditional translation using a statement and target language. Another task reads that output text and detects the language.\n",
        "\n",
        "This workflow uses a sequences pipeline. The sequences pipeline loads a Hugging Face sequence to sequence model for inference, in this case [FLAN-T5](https://huggingface.co/google/flan-t5-base). The sequences pipeline takes a prompt as input and outputs the model inference result.\n",
        "\n",
        "It's important to note that a pipeline is simply a callable function. It can easily be replaced with a call to an external API."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OUc9gqTyAYnm",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "83300311-736c-47c8-bc16-ec0303274054"
      },
      "source": [
        "from txtai.pipeline import Sequences\n",
        "from txtai.workflow import Workflow, TemplateTask\n",
        "\n",
        "# Create sequences pipeline\n",
        "sequences = Sequences(\"google/flan-t5-large\")\n",
        "\n",
        "# Define workflow or chaining of tasks together.\n",
        "workflow = Workflow([\n",
        "    TemplateTask(\n",
        "        template=\"Translate '{statement}' to {language} if it's English\",\n",
        "        action=sequences\n",
        "    ),\n",
        "    TemplateTask(\n",
        "        template=\"What language is the following text? {text}\",\n",
        "        action=sequences\n",
        "    )\n",
        "])\n",
        "\n",
        "inputs = [\n",
        "    {\"statement\": \"Hello, how are you\", \"language\": \"French\"},\n",
        "    {\"statement\": \"Hallo, wie geht's dir\", \"language\": \"French\"}\n",
        "]\n",
        "\n",
        "print(list(workflow(inputs)))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "['French', 'German']\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's recap what happened here. The first workflow task conditionally translates text to a language if it's English.\n",
        "\n",
        "The first statement is `Hello, how are you` with a target language of French. So the statement is translated to French.\n",
        "\n",
        "The second statement is German, so it's not converted to French.\n",
        "\n",
        "The next step asks the model what the language is and it correctly prints `French` and `German`."
      ],
      "metadata": {
        "id": "_zz4Do8BV-Lk"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Prompt Workflow as YAML\n",
        "\n",
        "The same workflow above can be created with YAML configuration."
      ],
      "metadata": {
        "id": "iXDAKP4CX0W9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%writefile workflow.yml\n",
        "\n",
        "sequences:\n",
        "  path: google/flan-t5-large\n",
        "\n",
        "workflow:\n",
        "  chain:\n",
        "    tasks:\n",
        "      - task: template\n",
        "        template: Translate '{statement}' to {language} if it's English\n",
        "        action: sequences\n",
        "      - task: template\n",
        "        template: What language is the following text? {text}\n",
        "        action: sequences"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GwV5A9xRYtYs",
        "outputId": "ffe6ee65-95a7-46c6-e6b9-5324eab26ca8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Writing workflow.yml\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from txtai import Application\n",
        "\n",
        "app = Application(\"workflow.yml\")\n",
        "print(list(app.workflow(\"chain\", inputs)))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dr7Lv5S5X98e",
        "outputId": "d6ac0427-671d-4525-aa21-664430109af3"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "['French', 'German']\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "As expected, the same result! This is a matter of preference on how you want to create a workflow. One advantage of YAML workflows is that an API can easily be created from the workflow file."
      ],
      "metadata": {
        "id": "EGqiV45fYVse"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Prompt Workflow via an API call\n",
        "\n",
        "Let's say you want the workflow to be available via an API call. Well good news, txtai has a built in API mechanism using FastAPI. "
      ],
      "metadata": {
        "id": "9PqMU0bNYinf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Start an API service\n",
        "!CONFIG=workflow.yml nohup uvicorn \"txtai.api:app\" &> api.log &\n",
        "!sleep 60"
      ],
      "metadata": {
        "id": "vDxQj1ZIYsz3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import requests\n",
        "\n",
        "# Run API request\n",
        "requests.post(\"http://localhost:8000/workflow\", json={\"name\": \"chain\", \"elements\": inputs}).json()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "R1o08SVtZW7h",
        "outputId": "99875acd-18a8-4c2c-ead3-cb6975a4b2d2"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "['French', 'German']"
            ]
          },
          "metadata": {},
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Just like the previous steps, except through an API call. Let's run via cURL for good measure."
      ],
      "metadata": {
        "id": "B88mCrGFl5W-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%bash\n",
        "\n",
        "curl -s -X POST \"http://localhost:8000/workflow\" \\\n",
        "     -H \"Content-Type: application/json\" \\\n",
        "     --data @- << EOF\n",
        "{\n",
        "  \"name\": \"chain\",\n",
        "  \"elements\": [\n",
        "    {\"statement\": \"Hello, how are you\", \"language\": \"French\"},\n",
        "    {\"statement\": \"Hallo, wie geht's dir\", \"language\": \"French\"}\n",
        "  ]\n",
        "}\n",
        "EOF"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hRUyh0cQl_P2",
        "outputId": "9db8481d-0b6e-4a31-bdf6-5443df5f768a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[\"French\",\"German\"]"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "One last time, the same output is shown.\n",
        "\n",
        "If your primary development environment isn't Python, txtai does have API bindings for [JavaScript](https://github.com/neuml/txtai.js), [Rust](https://github.com/neuml/txtai.rs), [Go](https://github.com/neuml/txtai.go) and [Java](https://github.com/neuml/txtai.java).\n",
        "\n",
        "More information on the API is available [here](https://neuml.github.io/txtai/api/)."
      ],
      "metadata": {
        "id": "W0zL93WPoaCo"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Conversational chain\n",
        "\n",
        "Conversational search is another big area of focus in 2023. [txtchat](https://github.com/neuml/txtchat) is a framework for building conversational search applications. It relies heavily on txtai. Let's see a conversational example."
      ],
      "metadata": {
        "id": "q9WiFG6fpzw5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%writefile search.yml\n",
        "\n",
        "writable: false\n",
        "cloud:\n",
        "  provider: huggingface-hub\n",
        "  container: neuml/txtai-intro\n",
        "\n",
        "rag:\n",
        "  path: google/flan-t5-large\n",
        "  output: reference\n",
        "\n",
        "workflow:\n",
        "  search:\n",
        "    tasks:\n",
        "      - task: rag\n",
        "        template: |\n",
        "          Answer the following question using only the context below. Give a detailed answer.\n",
        "          Say 'I don't have data on that' when the question can't be answered.\n",
        "          Question: {text}\n",
        "          Context: \n",
        "        action: rag\n",
        "      - task: template\n",
        "        template: \"{answer}\\n\\nReference: {reference}\"\n",
        "        rules:\n",
        "          answer: I don't have data on that"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rM3Y551LqF-J",
        "outputId": "85623785-c15f-4996-9460-0644f69cf5bf"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Overwriting search.yml\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "app = Application(\"search.yml\")\n",
        "print(list(app.workflow(\"search\", [\"Tell me something about North America\"])))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1Elb8JANqpwX",
        "outputId": "b1f1ffa1-6c47-4d90-b6f1-8098d4dc45f8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[\"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg\\n\\nReference: 1\"]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The first thing the code above does is run an embeddings search to build a conversational context. That context is then used to build a prompt and inference is run against the LLM. \n",
        "\n",
        "The next task formats the outputs with a reference to the best matching record. In this case, it's only an id of 1. But this can be much more useful if the id is a URL or there is logic to format the id back to a unique reference string.\n",
        "\n",
        "The [txtchat](https://github.com/neuml/txtchat) project has much more on this, check it out!"
      ],
      "metadata": {
        "id": "4r49V4c9s5nf"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Wrapping up\n",
        "\n",
        "This notebook covered how to build prompt templates and task chains through a series of results. txtai has long had a robust and efficient workflow framework for connecting models together. This can be small and simple models and/or prompting with large models. Go ahead and give it a try!"
      ],
      "metadata": {
        "id": "KqfvCXp2B3li"
      }
    }
  ]
}